from visual_localization import VisualLocalizer

from tornado.web import RequestHandler
import tornado.web
import tornado.httpserver
import tornado.ioloop
import numpy as np
from transforms3d.quaternions import *
from transforms3d.euler import *
import cv2
import argparse
import sys

class VisualLocateHandler(RequestHandler):
    def prepare(self):
        # solve CORS problem
        self.set_header('Access-Control-Allow-Origin', '*')  
        self.set_header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS')  
        self.set_header('Access-Control-Allow-Headers', 'Content-Type')
        self.set_header('Content-Type', 'application/json')
        super().prepare()

    def options(self):
        # solve CORS problem
        self.set_status(204)
        self.finish()
    
    async def post(self):
        # busy check and set
        if self.application.ctx['busy']:
            self.set_status(503)
            self.write({'error': 'server is busy'})
            return
        else:
            self.application.ctx['busy'] = True
        
        # get image from request
        image = None
        try:
            if 'image' not in self.request.files or len(self.request.files['image']) == 0:
                self.set_status(400)
                self.write({'error': 'no image in request'})
                self.application.ctx['busy'] = False
                return
            # get image from request
            image = self.request.files['image'][0]['body']
            image = np.frombuffer(image, np.uint8)
            image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        except:
            self.set_status(400)
            self.write({'error': 'invalid image'})
            self.application.ctx['busy'] = False
            return

        # localize
        localizer = self.application.ctx['localizer']
        result = await localizer.localize(image)
        if not result.success:
            self.set_status(500)
            self.write({'error': 'localization failed'})
            self.application.ctx['busy'] = False
            return
        
        R = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
        ROT = quat2mat(result.quaternion)
        print(result.quaternion)
        inv_rot = np.linalg.inv(ROT)
        result.translation = -inv_rot @ result.translation
        result.translation[2] = result.translation[2] + 0.625
        q = mat2quat(inv_rot @ R)
        result.quaternion = q
        print(result.translation, inv_rot)
        # return result
        response = {
            'success': True,
            'translation': result.translation.tolist(),
            'quaternion': result.quaternion.tolist(),
            'inlier_count': result.inlier_count,
            'reprojection_error': result.reprojection_error,
            'execution_time': result.execution_time,
            'confidence': result.confidence
        }
        
        # Add additional info if available
        if result.additional_info:
            response['additional_info'] = result.additional_info
        
        self.write(response)
        self.application.ctx['busy'] = False
        
class MyApplication(tornado.web.Application):
    def __init__(self, localizer):
        self.ctx = {
            'busy': False,
            'localizer': localizer
        }

        tornado.web.Application.__init__(self, [
            (r'/visual/locate', VisualLocateHandler)
        ])

if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Visual Localization HTTP Server')
    parser.add_argument('--reconstruction_path', type=str, required=True,
                        help='Path to COLMAP reconstruction (SfM model)')
    parser.add_argument('--db_global_features', type=str, required=True,
                        help='Path to database global features (NetVLAD) H5 file')
    parser.add_argument('--db_local_features', type=str, required=True,
                        help='Path to database local features (SuperPoint) H5 file')
    parser.add_argument('--top_k', type=int, default=20,
                        help='Number of top similar images to retrieve (default: 20)')
    parser.add_argument('--port', type=int, default=8049,
                        help='HTTP server port (default: 8049)')
    parser.add_argument('--host', type=str, default='0.0.0.0',
                        help='HTTP server host (default: 0.0.0.0)')
    
    args = parser.parse_args()
    
    print("="*80)
    print("Visual Localization HTTP Server")
    print("="*80)
    print(f"Reconstruction path:    {args.reconstruction_path}")
    print(f"Global features:        {args.db_global_features}")
    print(f"Local features:         {args.db_local_features}")
    print(f"Top K:                  {args.top_k}")
    print(f"Server address:         http://{args.host}:{args.port}")
    print("="*80)
    
    # Initialize visual localizer
    print("\n[Initialization] Loading models and data...")
    localizer = VisualLocalizer()
    
    try:
        success = localizer.initialize(
            reconstruction_path=args.reconstruction_path,
            db_global_features_path=args.db_global_features,
            db_local_features_path=args.db_local_features,
            top_k=args.top_k
        )
        
        if not success:
            print("✗ Visual localizer initialization failed!")
            sys.exit(1)
        
        print("✓ Visual localizer initialized successfully!")
        
    except Exception as e:
        print(f"✗ Error during initialization: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    
    # Initialize HTTP server
    print(f"\n[Server] Starting HTTP server on {args.host}:{args.port}...")
    app = MyApplication(localizer)
    web_server = tornado.httpserver.HTTPServer(app)
    web_server.listen(args.port, address=args.host)
    
    print("✓ Server is running!")
    print(f"✓ Endpoint: POST http://{args.host}:{args.port}/visual/locate")
    print("\nPress Ctrl+C to stop the server.\n")
    print("="*80)
    
    try:
        tornado.ioloop.IOLoop.current().start()
    except KeyboardInterrupt:
        print("\n\n" + "="*80)
        print("Server stopped by user.")
        print("="*80)
